import torch
from utils.torch_utils import select_device, time_sync
from models.common import DetectMultiBackend
from utils.general import (LOGGER, check_img_size, cv2, non_max_suppression, scale_coords)
from utils.dataloaders import LoadImages
from utils.plots import Annotator, colors
import numpy as np
import win32gui
import win32con
import win32ui
from PyQt5.QtWidgets import QApplication
import sys
import pyautogui
from PyQt5.QtWidgets import *


class SmartMa():
    def __init__(self):
        # 顺子
        self.straight_arr = [[], [], []]
        # 刻子
        self.carving_arr = [[], [], [], []]
        # 可能成为顺子
        self.maybe_straight_arr = [[], [], []]
        # 可能成为刻子
        self.maybe_carving_arr = [[], [], [], []]
        # 被挑选后剩余可打的牌
        self.surplus_arr = [[], [], [], []]
        # 手牌, 0万 1条 2筒 3字牌(0东 1南 2西 3北 4中 5发 6白)
        self.current_arr = [[3, 5], [2, 3, 4, 6], [1, 1, 1, 2], [5, 5, 5]]
        self.device = "cpu"
        self.model_path = "best.pt"
        self.model = self.model_load(weights=self.model_path, device=self.device)

    # 检测
    def detect_img(self, source):
        model = self.model
        imgsz = [640, 640]
        conf_thres = 0.50  # 置信度阈值
        iou_thres = 0.45
        max_det = 15  # 最大检测数量
        view_img = False
        save_crop = False
        nosave = False
        classes = None
        agnostic_nms = False
        augment = False
        visualize = False
        line_thickness = 1  # 线厚度 px
        hide_labels = False  # 隐藏标签
        hide_conf = False  # 隐藏置信度
        half = False
        print(source)
        if source == "":
            print("图片不存在")
        else:
            source = str(source)
            device = select_device(self.device)
            stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
            imgsz = check_img_size(imgsz, s=stride)  # 图片尺寸检查
            save_img = not nosave and not source.endswith('.txt')
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
            dt, seen = [0.0, 0.0, 0.0], 0
            for path, im, im0s, vid_cap, s in dataset:
                t1 = time_sync()
                im = torch.from_numpy(im).to(device)
                im = im.half() if half else im.float()
                im /= 255
                if len(im.shape) == 3:
                    im = im[None]
                t2 = time_sync()
                dt[0] += t2 - t1
                pred = model(im, augment=augment, visualize=visualize)
                t3 = time_sync()
                dt[1] += t3 - t2
                pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
                dt[2] += time_sync() - t3
                for i, det in enumerate(pred):
                    seen += 1
                    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
                    s += '%gx%g ' % im.shape[2:]
                    annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                    if len(det):
                        det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
                        for c in det[:, -1].unique():
                            n = (det[:, -1] == c).sum()
                            s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "

                        for *xyxy, conf, cls in reversed(det):
                            if save_img or save_crop or view_img:
                                c = int(cls)
                                label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                                print("data", xyxy, label)

                                annotator.box_label(xyxy, label, color=colors(c, True))

                    LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
                    im0 = annotator.result()
                    cv2.imwrite("./tmp/single_result.jpg", im0)

    # 模型初始化
    def model_load(self, weights="", device=''):
        device = select_device(device)
        model = DetectMultiBackend(weights, device=device, dnn=False)
        stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
        if pt:
            model.model.float()
        print("模型加载完成")
        return model

    # 寻找顺子
    def find_straight(self, cards):
        for idxi, i in enumerate(cards):
            cards_len = len(i)
            # 如果元素个数小于三则跳过
            if (cards_len < 3):
                continue
            for idxj, j in enumerate(i):
                next_card = i[idxj + 1]
                if ((idxj + 2) >= cards_len):
                    break
                next_next_card = i[idxj + 2]
                if ((idxj + 2) < cards_len and (next_card[0] - j[0]) == 1 and (next_next_card[0] - next_card[0]) == 1):
                    # 若满足则添加到对应数组
                    self.straight_arr[idxi].append((j, next_card, next_next_card))
                    i.pop(idxj)
                    i.pop(idxj)
                    i.pop(idxj)
                    # 添加后删除元素
        return cards

    # 寻找刻子
    def find_carving(self, cards):
        for idxi, i in enumerate(cards):
            cards_len = len(i)
            # 如果元素个数小于三则跳过
            if (cards_len < 3):
                continue
            for idxj, j in enumerate(i):
                next_card = i[idxj + 1]
                if ((idxj + 2) >= cards_len):
                    break
                next_next_card = i[idxj + 2]
                if ((idxj + 2) < cards_len and (next_card[0] == j[0]) and (next_next_card[0] == next_card[0])):
                    # 若满足则添加到对应数组
                    self.carving_arr[idxi].append((j, next_card, next_next_card))
                    i.pop(idxj)
                    i.pop(idxj)
                    i.pop(idxj)

        return cards

    # 寻找可能成为顺子
    def find_maybe_straight(self, cards):
        for idxi, i in enumerate(cards):
            cards_len = len(i)
            # 如果元素个数小于三则跳过
            if (cards_len < 2):
                continue
            for idxj, j in enumerate(i):
                if ((idxj + 1) >= cards_len):
                    break
                next_card = i[idxj + 1]
                if ((next_card[0] - j[0]) == 2):
                    self.maybe_straight_arr[idxi].append((j, next_card))
                    i.pop(idxj)
                    i.pop(idxj)
        return cards

    # 寻找可能成为刻子
    def find_maybe_carving(self, cards):
        for idxi, i in enumerate(cards):
            cards_len = len(i)
            # 如果元素个数小于三则跳过
            if (cards_len < 2):
                continue
            for idxj, j in enumerate(i):
                if ((idxj + 1) >= cards_len):
                    break
                next_card = i[idxj + 1]
                if (next_card[0] == j[0]):
                    self.maybe_carving_arr[idxi].append((j, next_card))
                    i.pop(idxj)
                    i.pop(idxj)
        return cards

    # 判断牌型
    def judge_status(self):
        sarr = [len(k) for k in self.straight_arr]
        slen = sum(sarr)
        print("当前手牌的顺子数量", slen)
        carr = [len(k) for k in self.carving_arr]
        clen = sum(carr)
        print("当前手牌的刻子数量", clen)
        darr = [len(k) for k in self.maybe_carving_arr]
        dlen = sum(darr)
        print("当前手牌的对子数量", dlen)
        if ((slen + clen) == 4 and dlen == 1):
            print("达成胡牌条件")

    def abandon(self):
        aarr = [len(k) for k in self.surplus_arr]
        aidx = [idx if i == 1 else -1 for idx, i in enumerate(aarr)]
        aidx.sort(reverse=True)
        if (aidx[0] != -1):
            return self.surplus_arr[aidx[0]]
            print("弃牌队列", self.surplus_arr, "弃牌索引：", aidx, "打出的牌：", self.surplus_arr[aidx[0]])
        else:
            return ()
            print("弃牌队列", self.surplus_arr, "弃牌索引：", aidx, "出错了啊铁子")

    def go(self):
        cards = self.find_straight(self.current_arr)
        print("找顺子后的手牌", cards, "顺子数组", self.straight_arr)
        cards = self.find_carving(cards)
        print("找刻子后的手牌", cards, "刻子数组", self.carving_arr)
        cards = self.find_maybe_straight(cards)
        print("找可能顺子后的手牌", cards, "可能顺子数组", self.maybe_straight_arr)
        cards = self.find_maybe_carving(cards)
        print("找可能刻子后的手牌", cards, "可能刻子数组", self.maybe_carving_arr)
        self.surplus_arr = cards
        print("可打的手牌", self.surplus_arr)


# 指定窗口截图
def winShot(roi, fmt="RGB"):
    left, top, w, h = roi
    # hwnd = win32gui.FindWindow(None, "pycharm")
    hwnd = win32gui.GetDesktopWindow()
    wDC = win32gui.GetWindowDC(hwnd)
    dcObj = win32ui.CreateDCFromHandle(wDC)
    cDC = dcObj.CreateCompatibleDC()
    dataBitMap = win32ui.CreateBitmap()

    dataBitMap.CreateCompatibleBitmap(dcObj, w, h)

    cDC.SelectObject(dataBitMap)
    cDC.BitBlt((0, 0), (w, h), dcObj, (left, top), win32con.SRCCOPY)
    signedIntsArray = dataBitMap.GetBitmapBits(True)
    img = np.frombuffer(signedIntsArray, dtype="uint8")
    img.shape = (h, w, 4)

    # Free Resources
    dcObj.DeleteDC()
    cDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, wDC)
    win32gui.DeleteObject(dataBitMap.GetHandle())

    if fmt == 'BGR':
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGBA2BGR)
    if fmt == 'RGB':
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGBA2RGB)
    else:
        raise ValueError('erro!')


def test():
    sm = SmartMa()
    sm.go()
    sm.judge_status()
    sm.abandon()


def test001(img_path, model):
    t1 = time_sync()
    model.detect_img(img_path)
    t2 = time_sync()
    print("识别耗时", t2 - t1)


def test003():
    img = pyautogui.screenshot()
    open_cv_image = np.array(img)
    open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
    # cv2.imwrite('image.jpg', open_cv_image)
    cv2.imshow("test", open_cv_image)



def test005():
    AI = SmartMa()
    AI.checkStatus()


if __name__ == '__main__':
    ai = SmartMa()
    test001(r"./tmp/img.png", ai)